package cn.johnyu;

import cn.johnyu.service.MilvusService;
import cn.johnyu.util.AliEmbeddingClient;
import cn.johnyu.util.MilvusUtil;
import io.milvus.client.MilvusServiceClient;

import java.util.Arrays;
import java.util.List;

public class EmbeddingMilvusOperate {
    private static List<String> mockTexts(){
        List<String> texts = Arrays.asList(
                "中华人民共和国刑法（第232条）: 故意杀人的，处死刑、无期徒刑或者十年以上有期徒刑；情节较轻的，处三年以上十年以下有期徒刑。",
                "中华人民共和国刑法（第264条）: 盗窃公私财物，数额较大的，或者多次盗窃、入户盗窃、携带凶器盗窃、扒窃的，处三年以下有期徒刑、拘役或者管制，并处或者单处罚金；数额巨大或者有其他严重情节的，处三年以上十年以下有期徒刑，并处罚金；数额特别巨大或者有其他特别严重情节的，处十年以上有期徒刑或者无期徒刑，并处罚金或者没收财产。",
                "中华人民共和国刑法（第133条）: 在道路上驾驶机动车，有下列情形之一的，处拘役，并处罚金：追逐竞驶、醉驾、严重超载超速、非法运输危险化学品等。",
                "中华人民共和国刑法（第382条）: 国家工作人员利用职务上的便利，侵吞、窃取、骗取或者以其他手段非法占有公共财物的，是贪污罪。",
                "中华人民共和国刑法（第266条）: 诈骗公私财物，数额较大的，处三年以下有期徒刑、拘役或者管制，并处或者单处罚金；数额巨大或者有其他严重情节的，处三年以上十年以下有期徒刑，并处罚金；数额特别巨大或者有其他特别严重情节的，处十年以上有期徒刑或者无期徒刑，并处罚金或者没收财产。",
                "中华人民共和国刑法（第234条）: 故意伤害他人身体的，处三年以下有期徒刑、拘役或者管制。致人重伤的，处三年以上十年以下有期徒刑；致人死亡或造成严重残疾的，处十年以上有期徒刑、无期徒刑或者死刑。",
                "中华人民共和国刑法（第130条）: 非法持有枪支、弹药、爆炸物的，构成犯罪的，依法追究刑事责任。",
                "中华人民共和国刑法（第310条）: 窝藏、包庇犯罪分子的，处三年以下有期徒刑、拘役或者管制；情节严重的，处三年以上七年以下有期徒刑。",
                "中华人民共和国刑法（第287条）: 利用计算机信息网络实施诈骗、盗窃、敲诈勒索的，依法追究刑事责任。",
                "中华人民共和国刑法（第299条）: 冒充国家机关工作人员招摇撞骗的，处三年以下有期徒刑、拘役、管制或者罚金。"
        );
        return texts;
    }
    public static void main(String[] args) throws Exception{
        String collectionName = "law_articles106";
        int dim=1024;
        String question="我使用爬虫技术盗取了某公司的敏感数据，并用于商用，应该触犯了哪条刑法";

        // 模拟文本数据,并生成向量列表
        List<String> texts=mockTexts();
        List<List<Float>> vectors = AliEmbeddingClient.embedding(texts);
        //获取问题的向量
        List<Float> queryVector = AliEmbeddingClient.embedding(Arrays.asList(question)).get(0);

        MilvusServiceClient milvusServiceClient = MilvusUtil.getClient("localhost", 19530);

        boolean exist= MilvusService.isCollectionExists(milvusServiceClient,collectionName,false);
        if (!exist) {
            int status = MilvusService.createCollection(milvusServiceClient, collectionName, dim);
            MilvusService.createIndex(milvusServiceClient, collectionName, "embedding");
        }
        //插入向量
        long count = MilvusService.insertData(milvusServiceClient, collectionName, texts, vectors);
//        //向量搜索
//        MilvusService.searchData(milvusServiceClient, collectionName, queryVector);

        MilvusUtil.close();
    }
}
